"""
WebSocket实时数据获取模块
提供实时K线和行情数据，替代REST API轮询
"""
import asyncio
import json
import time
import threading
from datetime import datetime
from typing import Callable, Dict, List, Optional
from collections import deque
import websocket

from src.utils.logger import get_logger

logger = get_logger(__name__)


class KlineBuffer:
    """K线缓冲区 - 存储最新的K线数据"""

    def __init__(self, max_size: int = 500):
        self.max_size = max_size
        self.klines: Dict[str, deque] = {}  # {symbol_timeframe: deque}

    def add_kline(self, symbol: str, timeframe: str, kline_data: Dict):
        """添加K线数据"""
        key = f"{symbol}_{timeframe}"

        if key not in self.klines:
            self.klines[key] = deque(maxlen=self.max_size)

        self.klines[key].append(kline_data)

    def get_latest_klines(self, symbol: str, timeframe: str, limit: int = 100) -> List[Dict]:
        """获取最新的K线数据"""
        key = f"{symbol}_{timeframe}"

        if key not in self.klines:
            return []

        klines = list(self.klines[key])
        return klines[-limit:] if len(klines) > limit else klines

    def get_latest_kline(self, symbol: str, timeframe: str) -> Optional[Dict]:
        """获取最新的单根K线"""
        key = f"{symbol}_{timeframe}"

        if key not in self.klines or len(self.klines[key]) == 0:
            return None

        return self.klines[key][-1]


class TickerBuffer:
    """行情数据缓冲区"""

    def __init__(self):
        self.tickers: Dict[str, Dict] = {}  # {symbol: ticker_data}

    def update_ticker(self, symbol: str, ticker_data: Dict):
        """更新行情数据"""
        self.tickers[symbol] = ticker_data

    def get_ticker(self, symbol: str) -> Optional[Dict]:
        """获取行情数据"""
        return self.tickers.get(symbol)


class OKXWebSocketClient:
    """OKX WebSocket 客户端"""

    # WebSocket端点
    LIVE_WS_URL = "wss://ws.okx.com:8443/ws/v5/public"
    SANDBOX_WS_URL = "wss://wspap.okx.com:8443/ws/v5/public"
    PRIVATE_WS_URL = "wss://ws.okx.com:8443/ws/v5/private"

    def __init__(self, api_key: str = "", secret_key: str = "", passphrase: str = "",
                 use_sandbox: bool = True, private: bool = False):
        """
        初始化WebSocket客户端

        Args:
            api_key: API密钥 (私有订阅需要)
            secret_key: 密钥 (私有订阅需要)
            passphrase: 密码短语 (私有订阅需要)
            use_sandbox: 是否使用沙盒环境
            private: 是否连接私有频道
        """
        self.api_key = api_key
        self.secret_key = secret_key
        self.passphrase = passphrase
        self.use_sandbox = use_sandbox
        self.private = private

        # 选择URL
        if private:
            self.ws_url = self.PRIVATE_WS_URL
        else:
            self.ws_url = self.SANDBOX_WS_URL if use_sandbox else self.LIVE_WS_URL

        # WebSocket连接
        self.ws: Optional[websocket.WebSocketApp] = None
        self.is_connected = False

        # 数据缓冲区
        self.kline_buffer = KlineBuffer()
        self.ticker_buffer = TickerBuffer()

        # 订阅配置
        self.subscriptions: List[Dict] = []

        # 回调函数
        self.on_kline_callbacks: List[Callable] = []
        self.on_ticker_callbacks: List[Callable] = []
        self.on_error_callbacks: List[Callable] = []

        # 线程控制
        self.ws_thread: Optional[threading.Thread] = None
        self.running = False

    def add_kline_callback(self, callback: Callable):
        """添加K线数据回调"""
        self.on_kline_callbacks.append(callback)

    def add_ticker_callback(self, callback: Callable):
        """添加行情数据回调"""
        self.on_ticker_callbacks.append(callback)

    def add_error_callback(self, callback: Callable):
        """添加错误回调"""
        self.on_error_callbacks.append(callback)

    def subscribe_klines(self, symbol: str, timeframes: List[str]):
        """
        订阅K线数据

        Args:
            symbol: 交易对 (如 BTC-USDT)
            timeframes: 时间框架列表 (如 ['1m', '5m', '1H'])
        """
        for timeframe in timeframes:
            sub = {
                "op": "subscribe",
                "args": [
                    {
                        "channel": f"candle{timeframe}",
                        "instId": symbol
                    }
                ]
            }
            self.subscriptions.append(sub)
            logger.info(f"订阅K线: {symbol} {timeframe}")

    def subscribe_ticker(self, symbols: List[str]):
        """
        订阅行情数据

        Args:
            symbols: 交易对列表
        """
        for symbol in symbols:
            sub = {
                "op": "subscribe",
                "args": [
                    {
                        "channel": "tickers",
                        "instId": symbol
                    }
                ]
            }
            self.subscriptions.append(sub)
            logger.info(f"订阅行情: {symbol}")

    def _on_message(self, ws, message: str):
        """处理WebSocket消息"""
        try:
            data = json.loads(message)

            # 处理K线数据
            if "arg" in data and "channel" in data["arg"]:
                channel = data["arg"]["channel"]

                if "candle" in channel:
                    self._handle_kline_data(data)
                elif "tickers" in channel:
                    self._handle_ticker_data(data)

        except Exception as e:
            logger.error(f"处理WebSocket消息异常: {e}")
            for callback in self.on_error_callbacks:
                callback(e)

    def _handle_kline_data(self, data: Dict):
        """处理K线数据"""
        try:
            if "data" not in data or len(data["data"]) == 0:
                return

            symbol = data["arg"]["instId"]
            timeframe = data["arg"]["channel"].replace("candle", "")

            for kline in data["data"]:
                kline_obj = {
                    'timestamp': int(kline[0]),
                    'open': float(kline[1]),
                    'high': float(kline[2]),
                    'low': float(kline[3]),
                    'close': float(kline[4]),
                    'volume': float(kline[5]),
                    'quote_volume': float(kline[6]),
                    'confirm': kline[8],  # K线是否确认
                }

                # 存入缓冲区
                self.kline_buffer.add_kline(symbol, timeframe, kline_obj)

                # 触发回调
                for callback in self.on_kline_callbacks:
                    try:
                        callback(symbol, timeframe, kline_obj)
                    except Exception as e:
                        logger.error(f"K线回调异常: {e}")

        except Exception as e:
            logger.error(f"处理K线数据异常: {e}")

    def _handle_ticker_data(self, data: Dict):
        """处理行情数据"""
        try:
            if "data" not in data or len(data["data"]) == 0:
                return

            for ticker in data["data"]:
                ticker_obj = {
                    'symbol': ticker["instId"],
                    'last_price': float(ticker.get("last", 0)),
                    'bid': float(ticker.get("bidPx", 0)),
                    'ask': float(ticker.get("askPx", 0)),
                    'high_24h': float(ticker.get("high24h", 0)),
                    'low_24h': float(ticker.get("low24h", 0)),
                    'volume_24h': float(ticker.get("vol24h", 0)),
                    'change_24h': float(ticker.get("change24h", 0)),
                    'timestamp': int(ticker.get("ts", 0)),
                }

                # 存入缓冲区
                self.ticker_buffer.update_ticker(ticker_obj['symbol'], ticker_obj)

                # 触发回调
                for callback in self.on_ticker_callbacks:
                    try:
                        callback(ticker_obj)
                    except Exception as e:
                        logger.error(f"行情回调异常: {e}")

        except Exception as e:
            logger.error(f"处理行情数据异常: {e}")

    def _on_open(self, ws):
        """WebSocket连接建立"""
        logger.info("WebSocket连接已建立")
        self.is_connected = True

        # 发送订阅请求
        for sub in self.subscriptions:
            try:
                ws.send(json.dumps(sub))
                time.sleep(0.1)  # 避免请求过快
            except Exception as e:
                logger.error(f"订阅失败: {e}")

    def _on_close(self, ws, close_status_code, close_msg):
        """WebSocket连接关闭"""
        logger.warning(f"WebSocket连接已关闭: {close_status_code} {close_msg}")
        self.is_connected = False

    def _on_error(self, ws, error):
        """WebSocket错误"""
        logger.error(f"WebSocket错误: {error}")
        self.is_connected = False

        for callback in self.on_error_callbacks:
            callback(error)

    def connect(self):
        """连接WebSocket"""
        if self.is_connected:
            logger.warning("WebSocket已连接")
            return

        logger.info(f"正在连接WebSocket: {self.ws_url}")

        # 禁用WebSocket的日志
        websocket.enableTrace(False)

        self.ws = websocket.WebSocketApp(
            self.ws_url,
            on_message=self._on_message,
            on_open=self._on_open,
            on_close=self._on_close,
            on_error=self._on_error,
        )

        # 在单独的线程中运行WebSocket
        self.running = True
        self.ws_thread = threading.Thread(target=self._run_ws, daemon=True)
        self.ws_thread.start()

    def _run_ws(self):
        """运行WebSocket连接"""
        try:
            self.ws.run_forever()
        except Exception as e:
            logger.error(f"WebSocket运行异常: {e}")
            self.is_connected = False
            self.running = False

    def disconnect(self):
        """断开WebSocket连接"""
        logger.info("断开WebSocket连接")
        self.running = False

        if self.ws:
            self.ws.close()

    def wait_connection(self, timeout: int = 30):
        """等待连接建立"""
        start_time = time.time()

        while not self.is_connected:
            if time.time() - start_time > timeout:
                raise TimeoutError("WebSocket连接超时")

            time.sleep(0.1)

        logger.info("WebSocket连接已就绪")

    def get_klines(self, symbol: str, timeframe: str, limit: int = 100) -> List[Dict]:
        """获取K线数据"""
        return self.kline_buffer.get_latest_klines(symbol, timeframe, limit)

    def get_latest_kline(self, symbol: str, timeframe: str) -> Optional[Dict]:
        """获取最新K线"""
        return self.kline_buffer.get_latest_kline(symbol, timeframe)

    def get_ticker(self, symbol: str) -> Optional[Dict]:
        """获取行情数据"""
        return self.ticker_buffer.get_ticker(symbol)


class RealtimeDataManager:
    """实时数据管理器 - 统一管理WebSocket连接"""

    def __init__(self, api_key: str = "", secret_key: str = "", passphrase: str = "",
                 use_sandbox: bool = True):
        self.api_key = api_key
        self.secret_key = secret_key
        self.passphrase = passphrase
        self.use_sandbox = use_sandbox

        # 公开数据连接
        self.public_ws: Optional[OKXWebSocketClient] = None
        # 私有数据连接
        self.private_ws: Optional[OKXWebSocketClient] = None

        self.data_callbacks: Dict[str, List[Callable]] = {
            'kline': [],
            'ticker': [],
            'error': [],
        }

    def add_kline_callback(self, callback: Callable):
        """添加K线回调"""
        self.data_callbacks['kline'].append(callback)

    def add_ticker_callback(self, callback: Callable):
        """添加行情回调"""
        self.data_callbacks['ticker'].append(callback)

    def add_error_callback(self, callback: Callable):
        """添加错误回调"""
        self.data_callbacks['error'].append(callback)

    def init_public_data(self, symbols: List[str], timeframes: List[str]):
        """初始化公开数据连接"""
        logger.info("初始化公开数据WebSocket连接")

        self.public_ws = OKXWebSocketClient(
            use_sandbox=self.use_sandbox,
            private=False
        )

        # 添加回调
        for callback in self.data_callbacks['kline']:
            self.public_ws.add_kline_callback(callback)

        for callback in self.data_callbacks['ticker']:
            self.public_ws.add_ticker_callback(callback)

        for callback in self.data_callbacks['error']:
            self.public_ws.add_error_callback(callback)

        # 订阅数据
        for symbol in symbols:
            self.public_ws.subscribe_klines(symbol, timeframes)
            self.public_ws.subscribe_ticker([symbol])

        # 连接
        self.public_ws.connect()

    def init_private_data(self):
        """初始化私有数据连接 (账户、订单等)"""
        if not self.api_key or not self.secret_key:
            logger.warning("未配置API密钥，跳过私有数据连接")
            return

        logger.info("初始化私有数据WebSocket连接")

        self.private_ws = OKXWebSocketClient(
            api_key=self.api_key,
            secret_key=self.secret_key,
            passphrase=self.passphrase,
            use_sandbox=self.use_sandbox,
            private=True
        )

        # 添加回调
        for callback in self.data_callbacks['error']:
            self.private_ws.add_error_callback(callback)

        # 连接
        self.private_ws.connect()

    def wait_ready(self, timeout: int = 30):
        """等待数据连接就绪"""
        if self.public_ws:
            self.public_ws.wait_connection(timeout)

        if self.private_ws:
            self.private_ws.wait_connection(timeout)

    def get_klines(self, symbol: str, timeframe: str, limit: int = 100) -> List[Dict]:
        """获取K线数据"""
        if not self.public_ws:
            return []

        return self.public_ws.get_klines(symbol, timeframe, limit)

    def get_latest_kline(self, symbol: str, timeframe: str) -> Optional[Dict]:
        """获取最新K线"""
        if not self.public_ws:
            return None

        return self.public_ws.get_latest_kline(symbol, timeframe)

    def get_ticker(self, symbol: str) -> Optional[Dict]:
        """获取行情数据"""
        if not self.public_ws:
            return None

        return self.public_ws.get_ticker(symbol)

    def disconnect(self):
        """断开所有连接"""
        if self.public_ws:
            self.public_ws.disconnect()

        if self.private_ws:
            self.private_ws.disconnect()
